#include "pch.h"
#include "TpFileManager.h"
#include "AsyncLibImpl.h"

namespace AsyncLib
{

CTpFileManager::CTpFileManager()
{
}

CTpFileManager::~CTpFileManager()
{
	assert(m_mapTasks.empty());
}

void CTpFileManager::Init()
{
}

void CTpFileManager::Stop()
{
	std::unique_lock lock(m_mutex);

	for (auto& [pTpIo, pTask] : m_mapTasks)
	{
		pTask->Cancel(FALSE);
	}
}

void CTpFileManager::Exit()
{
	for (auto& [pTpIo, pTask] : m_mapTasks)
	{
		pTask->Cancel(TRUE);
	}

	// They have been waited and closed outside
	m_mapTasks.clear();
}

std::shared_ptr<IFileTask> CTpFileManager::CreateTask(std::weak_ptr<CTpFileManager> pManager, HANDLE hFile)
{
	auto pTpIo = CreateThreadpoolIo(hFile, [](PTP_CALLBACK_INSTANCE, PVOID pContext, PVOID lpOverlapped, ULONG nIoResult, ULONG_PTR nBytesTransferred, PTP_IO pTpIo)
	{
		((CTpFileManager *)pContext)->OnCallback(pTpIo, lpOverlapped, (UINT)nBytesTransferred, nIoResult);
	}, this, g_asyncLib.GetCallbackEnv());

	assert(pTpIo);

	if (!pTpIo)
		return nullptr;

	auto pFileTask = std::make_shared<CTpFileTask>(std::move(pManager), hFile, pTpIo);
	std::unique_lock lock(m_mutex);
	m_mapTasks.emplace(pTpIo, pFileTask);
	return pFileTask;
}

void CTpFileManager::DeleteTask(PTP_IO pTpIo)
{
	std::unique_lock lock(m_mutex);

	if (g_asyncLib.IsRunning())
	{
		[[maybe_unused]] auto nCount = m_mapTasks.erase(pTpIo);
		assert(nCount == 1);
		CloseThreadpoolIo(pTpIo);
	}
}

void CTpFileManager::OnCallback(PTP_IO pTpIo, LPVOID lpOverlapped, UINT nBytesTransferred, DWORD dwError)
{
	g_asyncLib.SetInThreadPool();
	std::shared_lock lock(m_mutex);

	if (auto it = m_mapTasks.find(pTpIo); it != m_mapTasks.end() && g_asyncLib.IsRunning())
	{
		auto pFileTask = it->second;
		lock.unlock();
		pFileTask->DoCallback(lpOverlapped, nBytesTransferred, dwError);
	}
}

void CTpFileTask::DoCallback(LPVOID lpOverlapped, UINT nBytesTransferred, DWORD dwError)
{
	for (auto& node : m_caches)
	{
		if (&node.overlapped == lpOverlapped)
		{
			auto fnCallback = std::move(node.fnCallback);
			[[maybe_unused]] BOOL bUsing = InterlockedExchange(&node.bUsing, FALSE);
			assert(bUsing);

			if (fnCallback)
				fnCallback(nBytesTransferred, dwError);
			return;
		}
	}

	assert(false);
}

LPOVERLAPPED CTpFileTask::AllocCache(INT64 nOffset, std::function<void(UINT, DWORD)>&& fnCallback)
{
	assert(fnCallback);

	for (auto& node : m_caches)
	{
		if (!InterlockedExchange(&node.bUsing, TRUE))
		{
			node.fnCallback = std::move(fnCallback);

			LARGE_INTEGER liOffset{};
			liOffset.QuadPart = nOffset;
			node.overlapped.Offset = liOffset.LowPart;
			node.overlapped.OffsetHigh = liOffset.HighPart;
			return &node.overlapped;
		}
	}

	assert(false);

	if (fnCallback)
		fnCallback(0, ERROR_IO_INCOMPLETE);

	return nullptr;
}

void CTpFileTask::ExecuteIo(DWORD (*pfInvoke)(void *, LPOVERLAPPED), void *pfnIo, INT64 nOffset, std::function<void(UINT, DWORD)>&& fnCallback)
{
	if (auto pManager = m_pManager.lock())
	{
		std::shared_lock lock(pManager->GetSharedLock()); // For IsRunning() and m_pTpIo

		if (g_asyncLib.IsRunning())
		{
			auto lpOverlapped = AllocCache(nOffset, std::move(fnCallback));
			if (!lpOverlapped)
				return;

			StartThreadpoolIo(m_pTpIo);
			DWORD dwError = pfInvoke(pfnIo, lpOverlapped);

			if (dwError && dwError != ERROR_IO_PENDING)
			{
				CancelThreadpoolIo(m_pTpIo);
				lock.unlock();
				DoCallback(lpOverlapped, 0, dwError);
			}
		}
	}
}

template<typename Fn>
static DWORD Invoke(void *pfnIo, LPOVERLAPPED lpOverlapped)
{
	return (*(Fn *)pfnIo)(lpOverlapped);
}

void CTpFileTask::Read(void *lpBuffer, INT64 nOffset, UINT nBytesToRead, std::function<void(UINT, DWORD)>&& fnCallback)
{
	auto fnIo = [this, lpBuffer, nBytesToRead](LPOVERLAPPED lpOverlapped)
	{
		return ::ReadFile(m_hFile, lpBuffer, nBytesToRead, NULL, lpOverlapped) ? 0 : GetLastError();
	};

	ExecuteIo(Invoke<decltype(fnIo)>, &fnIo, nOffset, std::move(fnCallback));
}

void CTpFileTask::Write(const void *lpBuffer, INT64 nOffset, UINT nBytesToWrite, std::function<void(UINT, DWORD)>&& fnCallback)
{
	auto fnIo = [this, lpBuffer, nBytesToWrite](LPOVERLAPPED lpOverlapped)
	{
		return ::WriteFile(m_hFile, lpBuffer, nBytesToWrite, NULL, lpOverlapped) ? 0 : GetLastError();
	};

	ExecuteIo(Invoke<decltype(fnIo)>, &fnIo, nOffset, std::move(fnCallback));
}

void CTpFileTask::Close()
{
#ifdef _DEBUG
	for (const auto& node : m_caches)
	{
		assert(!node.bUsing);
	}
#endif

	if (auto pManager = m_pManager.lock())
		pManager->DeleteTask(m_pTpIo);
}

void CTpFileTask::Cancel(BOOL bExit)
{
	bool bComplete = true;

	for (auto& node : m_caches)
	{
		if (node.bUsing)
		{
			DWORD nBytesTransferred = 0;
			if (!GetOverlappedResult(m_hFile, &node.overlapped, &nBytesTransferred, bExit))
				bComplete = false;

			if (bExit)
			{
				auto fnCallback = std::move(node.fnCallback);
				node.bUsing = FALSE;
			}
		}
	}

	if (!bComplete && !bExit)
		CancelIoEx(m_hFile, NULL);
}

} // namespace AsyncLib
